"""
G1 robot controller with ONNX policy integration.
"""

import numpy as np
import torch
import onnxruntime as ort

from .base import BaseController
from ..core.articulation import ArticulationSubset
from ..utils.math_utils import quat_rotate_inverse
from ..utils.joint_mapping import JointMapping


class G1MoveBySpeedController(BaseController):
    """G1 Move By Speed Controller - 复制源代码实现"""

    # 完全复制源代码的关节定义
    joint_names_sim = [
            'left_hip_pitch_joint',
            'right_hip_pitch_joint',
            'waist_yaw_joint',
            'left_hip_roll_joint',
            'right_hip_roll_joint',
            'left_hip_yaw_joint',
            'right_hip_yaw_joint',
            'left_knee_joint',
            'right_knee_joint',
            'left_shoulder_pitch_joint',
            'right_shoulder_pitch_joint',
            'left_ankle_pitch_joint',
            'right_ankle_pitch_joint',
            'left_shoulder_roll_joint',
            'right_shoulder_roll_joint',
            'left_ankle_roll_joint',
            'right_ankle_roll_joint',
            'left_shoulder_yaw_joint',
            'right_shoulder_yaw_joint',
            'left_elbow_joint',
            'right_elbow_joint',
            'left_wrist_roll_joint',
            'right_wrist_roll_joint',
            'left_wrist_pitch_joint',
            'right_wrist_pitch_joint',
            'left_wrist_yaw_joint',
            'right_wrist_yaw_joint',
        ]

    joint_names_gym = [
            'left_hip_pitch_joint',
            'left_hip_roll_joint',
            'left_hip_yaw_joint',
            'left_knee_joint',
            'left_ankle_pitch_joint',
            'left_ankle_roll_joint',
            'right_hip_pitch_joint',
            'right_hip_roll_joint',
            'right_hip_yaw_joint',
            'right_knee_joint',
            'right_ankle_pitch_joint',
            'right_ankle_roll_joint',
            'waist_yaw_joint',
            'left_shoulder_pitch_joint',
            'left_shoulder_roll_joint',
            'left_shoulder_yaw_joint',
            'left_elbow_joint',
            'left_wrist_roll_joint',
            'left_wrist_pitch_joint',
            'left_wrist_yaw_joint',
            'right_shoulder_pitch_joint',
            'right_shoulder_roll_joint',
            'right_shoulder_yaw_joint',
            'right_elbow_joint',
            'right_wrist_roll_joint',
            'right_wrist_pitch_joint',
            'right_wrist_yaw_joint',
        ]

    def __init__(self, model_path, robot=None):
        super().__init__(robot)
        self.joint_mapping = JointMapping()

        # ONNX模型相关
        self.session = ort.InferenceSession(model_path)
        self.policy_input_obs_num = 546

        # ✅ 优化：添加joint_subset初始化
        self.joint_names = G1MoveBySpeedController.joint_names_sim
        self.joint_subset = ArticulationSubset(self.robot.articulation, self.joint_names)

        # 默认关节位置
        self.default_dof_pos = np.concatenate([
            np.array([-0.1, 0.0, 0.0, 0.3, -0.2, 0.0, -0.1, 0.0, 0.0, 0.3, -0.2, 0.0]),  # lower joints
            np.zeros(15),  # upper joints
        ])

        # 状态变量
        self._last_target_joint_positions = np.zeros(27)
        self._old_policy_obs = torch.zeros(self.policy_input_obs_num, dtype=torch.float32)
        self._apply_times_left = 0
        self.applied_joint_positions = self.default_dof_pos.copy()

        # ✅ Fast Fail: 直接获取底层_rigid_prim，完全按照原始源代码设计
        from isaacsim.core.utils.prims import get_prim_at_path  # 使用正确的导入路径
        from omni.isaac.core.prims import RigidPrim  # 使用正确的导入路径
        
        # 直接获取IMU prim并验证
        self._imu_prim = get_prim_at_path(self.robot._imu_in_torso)
        if not self._imu_prim.IsValid():
            raise RuntimeError(f"IMU prim {self.robot._imu_in_torso} 不存在")
        
        # ✅ 关键：直接使用_rigid_prim，就像原始源代码一样
        self._rigid_prim = RigidPrim(prim_path=self.robot._imu_in_torso)
        
        # 验证_rigid_prim是否正常工作
        try:
            # 测试get_angular_velocity()方法（就像原始源代码一样）
            test_ang_vel = self._rigid_prim.get_angular_velocity()
            if test_ang_vel is None:
                raise RuntimeError(f"IMU {self.robot._imu_in_torso} 无法获取角速度")
        except Exception as e:
            raise RuntimeError(f"IMU _rigid_prim初始化失败: {e}")

    def quat_rotate_inverse(self, q, v):
        """复制源代码的数学函数"""
        # 直接调用全局函数，避免重复实现
        return quat_rotate_inverse(q, v)

    def forward(self, forward_speed=0.0, lateral_speed=0.0, rotation_speed=0.0, height=0.74):
        """forward方法 - Fast Fail范式获取IMU数据"""
        if self._apply_times_left > 0:
            self._apply_times_left -= 1
            return self.joint_subset.make_articulation_action(
                joint_positions=self.applied_joint_positions, joint_velocities=None
            )

        # ✅ Fast Fail: 直接使用底层_rigid_prim，完全按照原始源代码模式
        # 使用_rigig_prim获取角速度（就像原始源代码一样）
        imu_ang_vel_w = self._rigid_prim.get_angular_velocity()
        
        # 使用_rigid_prim获取位姿（就像原始源代码一样）
        imu_position, imu_quaternion = self._rigid_prim.get_world_pose()
        imu_quaternion = np.array(imu_quaternion, dtype=np.float32)

        # 调试输出
        print(f"🔍 IMU数据 (Fast Fail _rigid_prim API):")
        print(f"   四元数 (w,x,y,z): {imu_quaternion}")
        print(f"   角速度 (x,y,z): {imu_ang_vel_w}")

        imu_quat_w = torch.tensor(imu_quaternion, dtype=torch.float32).reshape(1, -1)
        imu_ang_vel_w_tensor = torch.tensor(imu_ang_vel_w, dtype=torch.float32).reshape(1, -1)

        # 角速度转换到机体坐标系
        imu_ang_vel = np.array(quat_rotate_inverse(imu_quat_w, imu_ang_vel_w_tensor).reshape(-1))

        # 重力投影计算
        projected_gravity_tensor = torch.tensor([[0.0, 0.0, -1.0]], dtype=torch.float32)
        projected_gravity = np.array(quat_rotate_inverse(imu_quat_w, projected_gravity_tensor).reshape(-1))

        # 获取关节状态
        joint_pos = self.joint_mapping.sim2gym(self.joint_subset.get_joint_positions())
        joint_vel = self.joint_mapping.sim2gym(self.joint_subset.get_joint_velocities())

        # 计算关节位置偏差
        joint_pos -= self.default_dof_pos

        # 构建观测向量
        tracking_command = np.array([forward_speed, lateral_speed, rotation_speed], dtype=np.float32)
        current_obs = np.concatenate([
            tracking_command * np.array([2.0, 2.0, 0.25], dtype=np.float32),  # dim = 3
            np.array([height], dtype=np.float32),                            # dim = 1
            (imu_ang_vel * 0.25).astype(np.float32),                        # dim = 3
            projected_gravity.astype(np.float32),                           # dim = 3
            joint_pos.astype(np.float32),                                    # dim = 27
            (joint_vel * 0.05).astype(np.float32),                         # dim = 27
            self._last_target_joint_positions.astype(np.float32),           # dim = 27
        ])
        current_obs = np.clip(current_obs, -100.0, 100.0)

        # print(current_obs)

        # exit()

        # 构建历史观测
        policy_obs = torch.cat([self._old_policy_obs[91:], torch.tensor(current_obs, dtype=torch.float32)])
        self._old_policy_obs = policy_obs
        policy_obs = policy_obs.reshape(1, self.policy_input_obs_num)

        # ONNX推理
        input_name = self.session.get_inputs()[0].name
        input_data = policy_obs.numpy().astype(np.float32)
        target_lower_joint_positions = self.session.run(None, {input_name: input_data})[0][0]

        # 构建完整目标位置
        target_upper_joint_pos = np.zeros(15)
        self._last_target_joint_positions = np.clip(
            np.concatenate([target_lower_joint_positions, target_upper_joint_pos]), -100.0, 100.0
        )

        # 应用动作缩放和偏移
        gym_applied_joint_positions = self._last_target_joint_positions * 0.25 + self.default_dof_pos
        self.applied_joint_positions = self.joint_mapping.gym2sim(gym_applied_joint_positions)
        self._apply_times_left = 3

        return self.joint_subset.make_articulation_action(
            joint_positions=self.applied_joint_positions, joint_velocities=None
        )

    def action_to_control(self, action):
        """实现BaseController接口"""
        if len(action) < 3:
            raise ValueError("action must contain at least 3 elements: [forward_speed, lateral_speed, rotation_speed]")

        forward_speed = action[0]
        lateral_speed = action[1]
        rotation_speed = action[2]
        height = 0.74 if len(action) <= 3 else action[3]

        return self.forward(
            forward_speed=forward_speed,
            lateral_speed=lateral_speed,
            rotation_speed=rotation_speed,
            height=height
        )